﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Kokomo.Mirror.Ia32;
using Kokomo.Lexical.Expressions;
using Kokomo.Lexical.Coding;
using Kokomo.PeInspector.Reflection.Disassembly;

namespace Kokomo.PeInspector.Native
{
	class NativeDecompilationContext : DecompilationContext
	{
		public NativeDecompilationContext(Ia32Disassembly disassembly)
		{
			if (disassembly == null) throw new ArgumentNullException("disassembly");

			this._flags = new Dictionary<Eflags, NativeVariableLineage>();
			this._pendingFlags = new List<Eflags>();

			this._stack = new Stack<NativeVariableLineage>();

			this.Statements = new List<Statement>();

			this._registers = new Dictionary<Register, NativeVariableLineage>();
			this.ExternalDependencies = new List<NativeDependency>();
			this._setRegisters = new List<Register>();
			this._pendingRegisters = new List<Register>();
			this._memoryVariables = new Dictionary<long, VariableReferenceExpression>();

			this._registerVariables = new Dictionary<Register, VariableReferenceExpression>();
			this._flagVariables = new Dictionary<Eflags, VariableReferenceExpression>();

			this.Disassembly = disassembly;
			this._blockEnumerator = disassembly.Blocks.GetEnumerator();

			// Add arguments
			this._registerVariables.Add(Registers.RCX, new UnboundVariableReferenceExpression("arg0"));
			this._registerVariables.Add(Registers.RDX, new UnboundVariableReferenceExpression("arg1"));
			this._registerVariables.Add(Registers.R8, new UnboundVariableReferenceExpression("arg2"));
			this._registerVariables.Add(Registers.R9, new UnboundVariableReferenceExpression("arg3"));

			this._registerVariables.Add(Registers.RAX, new UnboundVariableReferenceExpression("loc0"));
			this._registerVariables.Add(Registers.R10, new UnboundVariableReferenceExpression("loc1"));
			this._registerVariables.Add(Registers.R11, new UnboundVariableReferenceExpression("loc2"));
		}

		#region Stack frame
		public bool IsStackFrameAllocated { get; private set; }
		public int StackFrameSize { get; private set; }

		public void AllocateStackFrame(int size)
		{
			//if (this.IsStackFrameAllocated) throw new InvalidOperationException("Stack frame already allocated");

			this.IsStackFrameAllocated = true;
			this.StackFrameSize += size;

			this.Statements.Add(new ExpressionStatement(new LocalAllocExpression(ConstantExpression.FromValue(size))));
		}
		public void DeallocateStackFrame(int size)
		{
			this.StackFrameSize -= size;
		}
		#endregion
		public bool IsStackPointer(Register register)
		{
			if (register == null) throw new ArgumentNullException("register");

			bool isStackPointer =
				(register.Type == RegisterType.GeneralPurpose)
				&& (register.Number == Registers.ESP.Number)
				;
			return isStackPointer;
		}

		#region Position
		public Ia32Disassembly Disassembly { get; private set; }
		private IEnumerator<Ia32DisassemblyBlock> _blockEnumerator;
		private IEnumerator<Ia32Instruction> _blockInstructionEnumerator;

		public Ia32Instruction CurrentInstruction { get; private set; }
		public long CurrentOffset { get { return this.CurrentInstruction.Offset; } }
		public bool End { get; private set; }
		public bool MoveNext()
		{
			if (this.End) return false;

			bool next = (this._blockInstructionEnumerator != null) && this._blockInstructionEnumerator.MoveNext();
			if (!next)
			{
				next = this._blockEnumerator.MoveNext();
				if (!next)
				{
					this.End = false;
				}
				else
				{
					this._blockInstructionEnumerator = this._blockEnumerator.Current.Instructions.GetEnumerator();
					this._blockInstructionEnumerator.MoveNext();
				}

				this.OnNextBlock();
			}

			this.CurrentInstruction = next
				? this._blockInstructionEnumerator.Current
				: null
				;

			return next;
		}

		private void OnNextBlock()
		{
			FlushFlagsAndRegisters();
		}

		private void FlushFlagsAndRegisters()
		{
			foreach (var flag in this._pendingFlags)
			{
				NativeVariableLineage lineage = this._flags[flag];
				Expression expression = lineage.Value;
				Expression variable = this.GetVariableForFlag(flag);
				this.Statements.Add(new AssignmentStatement(variable, expression));
				this._flags[flag] = new NativeVariableLineage(variable, lineage.InstructionOffset, lineage.InferredType, lineage.Dependencies);
			}
			this._pendingFlags.Clear();

			foreach (var register in this._pendingRegisters)
			{
				NativeVariableLineage lineage = this._registers[register];
				Expression expression = lineage.Value;
				Expression variable = this.GetVariableForRegister(register);
				this.Statements.Add(new AssignmentStatement(variable, expression));
				this._registers[register] = new NativeVariableLineage(variable, lineage.InstructionOffset, lineage.InferredType, lineage.Dependencies);
			}
			this._pendingRegisters.Clear();
		}
		#endregion

		#region Statements
		public IList<Statement> Statements { get; private set; }
		#endregion

		#region Flow
		public void OnBranching()
		{
			this.FlushFlagsAndRegisters();
		}
		#endregion

		#region Dependencies
		public IList<NativeDependency> ExternalDependencies { get; private set; }

		private IList<Register> _setRegisters;
		private IList<Register> _pendingRegisters;
		#endregion

		#region Registers
		private Dictionary<Register, NativeVariableLineage> _registers;
		#endregion

		#region Flags
		private Dictionary<Eflags, NativeVariableLineage> _flags;
		private IList<Eflags> _pendingFlags;

		public NativeVariableLineage GetFlag(Eflags flag)
		{
			NativeVariableLineage lineage;
			if (!this._flags.TryGetValue(flag, out lineage))
			{
				lineage = new NativeVariableLineage(
					new BoundVariableReferenceExpression(new WorkspaceVariable("initial" + flag.ToString(), typeof(bool))),
					this.CurrentOffset,
					typeof(bool),
					new NativeVariableLineage[0]
					);
				this._flags.Add(flag, lineage);
			}

			return lineage;
		}
		public void SetFlag(
			Eflags flag,
			Expression value,
			params NativeVariableLineage[] dependencies
			)
		{
			NativeVariableLineage lineage = new NativeVariableLineage(
				value,
				this.CurrentOffset,
				typeof(bool),
				dependencies
				);
			this._flags[flag] = lineage;
		}
		public void SetArithmeticFlags(NativeVariableLineage lineage)
		{
			if (lineage == null) throw new ArgumentNullException("lineage");

			Expression value = lineage.Value;
			this.SetFlag(Eflags.Carry, Expression.op_LogicalOr((value < ConstantExpression.ZeroInt), (value > ConstantExpression.FromValue(uint.MaxValue))), lineage);
			this.SetFlag(Eflags.Overflow, Expression.op_LogicalOr((value < ConstantExpression.FromValue(int.MinValue)), (value > ConstantExpression.FromValue(int.MaxValue))), lineage);
			this.SetFlag(Eflags.AuxCarry, new BoundFunctionCallExpression(IntrinsicFunctions.AuxCarryFunction, new Expression[] { value }), lineage);
			this.SetLogicalFlags(lineage);
		}
		public void SetLogicalFlags(NativeVariableLineage lineage)
		{
			if (lineage == null) throw new ArgumentNullException("lineage");

			Expression value = lineage.Value;
			this.SetFlag(Eflags.Zero, Expression.Equality(value, ConstantExpression.ZeroInt), lineage);
			this.SetFlag(Eflags.Sign, value < ConstantExpression.ZeroInt, lineage);
			this.SetFlag(Eflags.Parity, new BoundFunctionCallExpression(IntrinsicFunctions.ParityFunction, new Expression[] { value }), lineage);
		}
		#endregion

		#region Variables
		private Dictionary<long, VariableReferenceExpression> _memoryVariables;
		private VariableReferenceExpression GetVariableForAddress(long address)
		{
			VariableReferenceExpression variable;
			if (this._memoryVariables.TryGetValue(address, out variable))
				return variable;

			variable = new UnboundVariableReferenceExpression("var" + address.ToString("X"));
			this._memoryVariables.Add(address, variable);

			return variable;
		}

		private Dictionary<Register, VariableReferenceExpression> _registerVariables;
		private VariableReferenceExpression GetVariableForRegister(Register register)
		{
			if (register == null) throw new ArgumentNullException("register");

			VariableReferenceExpression variable;
			if (this._registerVariables.TryGetValue(register, out variable))
				return variable;

			variable = new UnboundVariableReferenceExpression("$" + register.Name);
			this._registerVariables.Add(register, variable);

			return variable;
		}

		private Dictionary<Eflags, VariableReferenceExpression> _flagVariables;
		private VariableReferenceExpression GetVariableForFlag(Eflags flag)
		{
			VariableReferenceExpression variable;
			if (this._flagVariables.TryGetValue(flag, out variable))
				return variable;

			variable = new UnboundVariableReferenceExpression("$" + flag.ToString());
			this._flagVariables.Add(flag, variable);

			return variable;
		}
		#endregion

		#region Operands
		public NativeVariableLineage OperandToLineage(Operand operand)
		{
			if (operand == null) throw new ArgumentNullException("operand");

			NativeVariableLineage lineage;

			Register register = operand as Register;
			if (register != null)
			{
				lineage = RegisterOperandToLineage(register);
			}
			else
			{
				IndirectOperand indirect = operand as IndirectOperand;
				if (indirect != null)
				{
					lineage = IndirectOperandToLineage(indirect);
				}
				else
				{
					IndexedOperand indexed = operand as IndexedOperand;
					if (indexed != null)
					{
						lineage = IndexedOperandToLineage(indexed);
					}
					else
					{
						ImmediateOperand immediate = operand as ImmediateOperand;
						if (immediate != null)
						{
							lineage = ImmediateOperandToLineage(immediate);
						}
						else
						{
							JumpTarget jumpTarget = operand as JumpTarget;
							if (jumpTarget != null)
							{
								lineage = JumpTargetOperandToLineage(jumpTarget);
							}
							else
							{
								throw new UnsupportedOperandException(operand);
							}
						}
					}
				}
			}

			return lineage;
		}
		private NativeVariableLineage JumpTargetOperandToLineage(JumpTarget jumpTarget)
		{
			return new NativeVariableLineage(ConstantExpression.FromValue(jumpTarget.TargetOffset), this.CurrentOffset, typeof(IntPtr), null);
		}
		private NativeVariableLineage RegisterOperandToLineage(Register register)
		{
			NativeVariableLineage lineage;
			if (!this._registers.TryGetValue(register, out lineage))
			{
				Type type = this.DetermineDefaultDataType(register);
				lineage = new NativeVariableLineage(
					new BoundVariableReferenceExpression(new WorkspaceVariable(register.Name + "Param", type)),
					this.CurrentOffset,
					type,
					new NativeVariableLineage[0]
					);
				this._registers.Add(register, lineage);

				if (!this._setRegisters.Contains(register))
				{
					this._setRegisters.Add(register);
					this.ExternalDependencies.Add(new NativeRegisterDependency(this.CurrentOffset, register));
				}
			}
			return lineage;
		}
		private NativeVariableLineage IndirectOperandToLineage(IndirectOperand indirect)
		{
			NativeVariableLineage lineage;
			//if (indirect.SegmentOverride == Registers.FS)
			//{
			//}
			//else
			{
				Expression baseExpression = this.OperandToLineage(indirect.BaseOperand).Value;

				Expression expression = baseExpression;
				if (indirect.Displacement > 0)
					expression = expression + ConstantExpression.FromValue(indirect.Displacement);
				else if (indirect.Displacement < 0)
					expression = expression - ConstantExpression.FromValue(-indirect.Displacement);

				expression = new IndirectExpression(expression);
				lineage = new NativeVariableLineage(expression, this.CurrentOffset, typeof(IntPtr), null);
			}
			return lineage;
		}
		private NativeVariableLineage IndexedOperandToLineage(IndexedOperand indexed)
		{
			NativeVariableLineage lineage;
			Expression baseExpression = this.OperandToLineage(indexed.BaseOperand).Value;

			Expression indexExpression = null;
			if (indexed.IndexOperand != null)
			{
				indexExpression = this.OperandToLineage(indexed.IndexOperand).Value;
				if (indexed.Scale != null)
					indexExpression *= ConstantExpression.FromValue(indexed.Scale);

				if (indexed.Displacement > 0)
					indexExpression += ConstantExpression.FromValue(indexed.Displacement);
				else if (indexed.Displacement < 0)
					indexExpression -= ConstantExpression.FromValue(-indexed.Displacement);
			}
			else if (indexed.Displacement != null)
			{
				indexExpression = ConstantExpression.FromValue(indexed.Displacement);
			}

			Expression expression;
			if (indexExpression != null)
				expression = new IndexerExpression(baseExpression, new Expression[] { indexExpression });
			else
				expression = new IndirectExpression(baseExpression);
			lineage = new NativeVariableLineage(expression, this.CurrentOffset, typeof(IntPtr), null);
			return lineage;
		}
		private NativeVariableLineage ImmediateOperandToLineage(ImmediateOperand immediate)
		{
			NativeVariableLineage lineage;
			lineage = new NativeVariableLineage(ConstantExpression.FromValue(immediate.Value), this.CurrentOffset, typeof(IntPtr), null);
			return lineage;
		}

		public NativeVariableLineage LoadOperand(Operand operand)
		{
			if (operand == null) throw new ArgumentNullException("operand");

			return this.OperandToLineage(operand);
		}
		public NativeVariableLineage SetOperand(Operand operand, Expression value, Type inferredType, params NativeVariableLineage[] dependencies)
		{
			if (value == null) throw new ArgumentNullException("value");
			if (inferredType == null) throw new ArgumentNullException("inferredType");

			NativeVariableLineage lineage = new NativeVariableLineage(
				value,
				this.CurrentOffset,
				inferredType,
				dependencies
				);

			this.SetOperand(
				operand,
				lineage
				);
			return lineage;
		}
		public void SetOperand(Operand operand, NativeVariableLineage lineage)
		{
			if (operand == null) throw new ArgumentNullException("operand");
			if (lineage == null) throw new ArgumentNullException("lineage");

			Register register = operand as Register;
			if (register != null)
			{
				this._registers[register] = lineage;
				if (!this._setRegisters.Contains(register))
					this._setRegisters.Add(register);
				if (!this._pendingRegisters.Contains(register))
					this._pendingRegisters.Add(register);
			}
			else
			{
				Expression left = this.OperandToLineage(operand).Value;

				this.Statements.Add(new AssignmentStatement(left, lineage.Value));
			}
		}
		#endregion

		private Type DetermineDefaultDataType(Register register)
		{
			if (register == null) throw new ArgumentNullException("register");

			Type defaultType;
			switch (register.Type)
			{
				case RegisterType.GeneralPurpose:
				case RegisterType.Control:
				case RegisterType.Segment:
				case RegisterType.Debug:
					defaultType = NativeDecompiler.GetTypeForOperandSize(register.Size, true);
					break;
				case RegisterType.Mmx:
				case RegisterType.Simd:
				case RegisterType.Fpu:
					defaultType = typeof(double);
					break;
				default:
					defaultType = typeof(int);
					break;
			}

			return defaultType;
		}

		#region Stack
		private Stack<NativeVariableLineage> _stack;

		public void Push(
			Expression value,
			Type inferredType,
			params NativeVariableLineage[] dependencies
			)
		{
			this.Push(new NativeVariableLineage(
				value,
				this.CurrentOffset,
				inferredType,
				dependencies
				));
		}
		public void Push(NativeVariableLineage lineage)
		{
			this._stack.Push(lineage);
		}
		public NativeVariableLineage Pop()
		{
			NativeVariableLineage lineage = this._stack.Pop();
			return lineage;
		}
		#endregion
	}
}
